#include "hypercube_test_helper.hpp"
#include <NTL/BasicThreadPool.h>
#include <omp.h>

INTEL_ITT_DOMAIN_CREATE(hmm_midk_domain, "test.hypercube.hmm.midk");
INTEL_ITT_STRING_HANDLE_CREATE(hmm_midk_handle, "HMM_MIDK");

typedef unsigned long ulong;

auto max_omp_threads = omp_get_max_threads();
ulong ntl_threads;
ulong hmm_threads;
ulong ntl_avail_threads;

bool bsgs;
bool no_encoding_reuse = false;

bool only_multi_thread;
bool only_single_thread;

bool cache_free = false;
bool minimal = false;

bool should_run_test_case(const std::string &test_case_name) {
  if (only_multi_thread && only_single_thread) {
    return test_case_name == "General";
  } else if (only_single_thread) {
    return test_case_name == "ST";
  } else if (only_multi_thread) {
    return test_case_name == "MT";
  } else {
    return test_case_name == "General";
  }
}

class MidkHMatMulTest : public MatMulTest {
  /* nothing */
};

#define MIDK_HMM_TEST_TEMPLATE(DIM_M, DIM_N, DIM_K)                            \
  class TEST_M##DIM_M##_N##DIM_N##_K##DIM_K : public MidkHMatMulTest {         \
  protected:                                                                   \
    void SetUp() override {                                                    \
      /* Get the test case name  */                                            \
      const ::testing::TestInfo *const test_info =                             \
          ::testing::UnitTest::GetInstance()->current_test_info();             \
      const std::string test_case_name = test_info->name();                    \
                                                                               \
      if (!should_run_test_case(test_case_name)) {                             \
        GTEST_SKIP();                                                          \
      }                                                                        \
      m_ = DIM_M;                                                              \
      n_ = DIM_N;                                                              \
      k_ = DIM_K;                                                              \
      const auto eh_mnk = shmm_engine::get_expected_mnk(m_, n_, k_);           \
      type_ = shmm_engine::parse_matrix_type(eh_mnk);                          \
      ASSERT_EQ(MATRIX_TYPE::MEDIUM_K, type_);                                 \
      hmm_status_.method(bsgs ? METHOD::BSGS : METHOD::NAIVE);                 \
                                                                               \
      cache_free_ = cache_free;                                                \
      minimal_m_ = minimal;                                                    \
      MidkHMatMulTest::SetUp();                                                \
    }                                                                          \
  };                                                                           \
                                                                               \
  TEST_F(TEST_M##DIM_M##_N##DIM_N##_K##DIM_K, General) {                       \
    auto threads = std::min(hmm_threads, static_cast<ulong>(DIM_K));           \
                                                                               \
    auto c_AB = hmmcc_->get_empty_ctxt();                                      \
    HELIB_NTIMER_START(_Encryption);                                           \
    auto c_A = hmmcc_->encrypt(A_);                                            \
    auto c_B = hmmcc_->encrypt(B_);                                            \
    HELIB_NTIMER_STOP(_Encryption);                                            \
    auto enc_time = _named_local_timer_Encryption.getTime() * 1000.0;          \
                                                                               \
    HELIB_NTIMER_START(_HMM1);                                                 \
    hmme_->cgemm(type_, c_AB, c_A, c_B);                                       \
    HELIB_NTIMER_STOP(_HMM1);                                                  \
                                                                               \
    NTL::mat_ZZ res;                                                           \
    hmmcc_->decrypt(res, c_AB);                                                \
    MatMulTest::check(AB_, res);                                               \
                                                                               \
    std::cout << "===================== HMM Report ======================\n";  \
    hmmcc_->report_time();                                                     \
    std::cout << "* Encryption Time            : " << enc_time << " ms\n";     \
    std::cout << "*" << std::endl;                                             \
    std::cout << "* ST-HMM Time (Outside)      : "                             \
              << _named_local_timer_HMM1.getTime() << " s\n";                  \
    std::cout << "* Single Thread              : " << std::endl;               \
    hmme_->report();                                                           \
                                                                               \
    c_AB.clear();                                                              \
    hmme_->status()->reset(true);                                              \
    hmme_->status()->threads(threads);                                         \
    HELIB_NTIMER_START(_HMM2);                                                 \
    hmme_->cgemm(type_, c_AB, c_A, c_B);                                       \
    HELIB_NTIMER_STOP(_HMM2);                                                  \
                                                                               \
    hmmcc_->decrypt(res, c_AB);                                                \
    MatMulTest::check(AB_, res);                                               \
                                                                               \
    std::cout << "*" << std::endl;                                             \
    std::cout << "* MT-HMM Time (Outside)      : "                             \
              << _named_local_timer_HMM2.getTime() << " s\n";                  \
    std::cout << "* Multiple Threads           : " << std::endl;               \
    std::cout << "  * max # of threads         : "                             \
              << hmme_->status()->threads() << std::endl;                      \
    std::cout << "  * max # of ntl threads     : " << ntl_avail_threads        \
              << std::endl;                                                    \
    std::cout << "  * max parallelism          : " << threads << std::endl;    \
    hmme_->report();                                                           \
    std::cout << "=======================================================\n";  \
  }                                                                            \
                                                                               \
  TEST_F(TEST_M##DIM_M##_N##DIM_N##_K##DIM_K, ST) {                            \
    auto c_AB = hmmcc_->get_empty_ctxt();                                      \
    HELIB_NTIMER_START(_Encryption);                                           \
    auto c_A = hmmcc_->encrypt(A_);                                            \
    auto c_B = hmmcc_->encrypt(B_);                                            \
    HELIB_NTIMER_STOP(_Encryption);                                            \
    auto enc_time = _named_local_timer_Encryption.getTime() * 1000.0;          \
                                                                               \
    INTEL_ITT_RESUME;                                                          \
    INTEL_ITT_TASK_BEGIN(hmm_midk_domain, hmm_midk_handle);                    \
    HELIB_NTIMER_START(_HMM);                                                  \
    hmme_->cgemm(type_, c_AB, c_A, c_B);                                       \
    HELIB_NTIMER_STOP(_HMM);                                                   \
    INTEL_ITT_TASK_END(hmm_midk_domain);                                       \
    INTEL_ITT_DETACH;                                                          \
                                                                               \
    HELIB_NTIMER_START(_Decryption);                                           \
    NTL::mat_ZZ res;                                                           \
    hmmcc_->decrypt(res, c_AB);                                                \
    std::cout << A_ << std::endl;                                              \
    std::cout << B_ << std::endl;                                              \
    std::cout << res << std::endl;                                             \
    HELIB_NTIMER_STOP(_Decryption);                                            \
    auto dec_time = _named_local_timer_Decryption.getTime() * 1000.0;          \
    MatMulTest::check(AB_, res);                                               \
                                                                               \
    std::cout << "===================== HMM Report ======================\n";  \
    hmmcc_->report_time();                                                     \
    std::cout << "* Encryption Time            : " << enc_time << " ms\n";     \
    std::cout << "* Decryption Time            : " << dec_time << " ms\n";     \
    std::cout << "* HMM Time (Outside)         : "                             \
              << _named_local_timer_HMM.getTime() << " s\n";                   \
    std::cout << "* Single Thread              : " << std::endl;               \
    hmme_->report();                                                           \
    std::cout << "=======================================================\n";  \
  }                                                                            \
                                                                               \
  TEST_F(TEST_M##DIM_M##_N##DIM_N##_K##DIM_K, MT) {                            \
    int threads = std::min(hmm_threads, static_cast<ulong>(DIM_K));            \
    hmme_->status()->threads(threads);                                         \
                                                                               \
    auto c_AB = hmmcc_->get_empty_ctxt();                                      \
    HELIB_NTIMER_START(_Encryption);                                           \
    auto c_A = hmmcc_->encrypt(A_);                                            \
    auto c_B = hmmcc_->encrypt(B_);                                            \
    HELIB_NTIMER_STOP(_Encryption);                                            \
    auto enc_time = _named_local_timer_Encryption.getTime() * 1000.0;          \
    if (no_encoding_reuse) {                                                   \
      hmme_->status()->set_encoding_reuse(false);                              \
    }                                                                          \
                                                                               \
    INTEL_ITT_RESUME;                                                          \
    INTEL_ITT_TASK_BEGIN(hmm_midk_domain, hmm_midk_handle);                    \
    HELIB_NTIMER_START(_HMM);                                                  \
    hmme_->cgemm(type_, c_AB, c_A, c_B);                                       \
    HELIB_NTIMER_STOP(_HMM);                                                   \
    INTEL_ITT_TASK_END(hmm_midk_domain);                                       \
    INTEL_ITT_DETACH;                                                          \
                                                                               \
    HELIB_NTIMER_START(_Decryption);                                           \
    NTL::mat_ZZ res;                                                           \
    hmmcc_->decrypt(res, c_AB);                                                \
    HELIB_NTIMER_STOP(_Decryption);                                            \
    auto dec_time = _named_local_timer_Decryption.getTime() * 1000.0;          \
    MatMulTest::check(AB_, res);                                               \
                                                                               \
    std::cout << "===================== HMM Report ======================\n";  \
    hmmcc_->report_time();                                                     \
    std::cout << "* Encryption Time            : " << enc_time << " ms\n";     \
    std::cout << "* Decryption Time            : " << dec_time << " ms\n";     \
    std::cout << "* HMM Time (Outside)         : "                             \
              << _named_local_timer_HMM.getTime() << " s\n";                   \
    std::cout << "* Multiple Threads           : " << std::endl;               \
    std::cout << "  * max # of threads         : "                             \
              << hmme_->status()->threads() << std::endl;                      \
    std::cout << "  * max # of ntl threads     : " << ntl_avail_threads        \
              << std::endl;                                                    \
    std::cout << "  * max parallelism          : " << threads << std::endl;    \
    hmme_->report();                                                           \
    std::cout << "=======================================================\n";  \
  }

MIDK_HMM_TEST_TEMPLATE(128, 10, 64)
MIDK_HMM_TEST_TEMPLATE(10, 128, 64)

int main(int argc, char **argv) {
  ::testing::InitGoogleTest(&argc, argv);

  INTEL_ITT_PAUSE; //! ITT - Stop

  command_line_parser parser(my_params);
  parser.add_option("-bsgs", "bsgs", &bsgs);
  parser.add_option("-ner", "no_encoding_reuse", &no_encoding_reuse);
  parser.add_option("-mt", "only_multi_thread", &only_multi_thread);
  parser.add_option("-ht", std::to_string(max_omp_threads), "hmm_threads",
                    &hmm_threads);
  parser.add_option("-nt", std::to_string(max_omp_threads), "ntl_threads",
                    &ntl_threads);
  parser.add_option("-st", "only_single_thread", &only_single_thread);
  parser.add_option("-nocache", "always_build_context", &cache_free);
  parser.add_option("-minimal", "find_minimal_m", &minimal);
  int ret = parser.parse(argc, argv);

  hmm_threads = std::min(hmm_threads, static_cast<ulong>(max_omp_threads));

  //* See https://libntl.org/doc/BasicThreadPool.cpp.html
  //! If we do not call `NTL::SetNumThreads(nt)`, then the
  //! `NTL::AvailableThreads()` will return 1 which indicates single-thread
  //
  //? why nt > 1 causes the Intel Vtune to crash
  NTL::SetNumThreads(
      std::min(ntl_threads, static_cast<ulong>(max_omp_threads)));
  ntl_avail_threads = NTL::AvailableThreads();

  if (ret == 0) {
    ret = RUN_ALL_TESTS();
  }

  return ret;
}